package stone926.mods.more_enchantments.mixins;

import net.minecraft.enchantment.EnchantmentHelper;
import net.minecraft.enchantment.Enchantments;
import net.minecraft.entity.EntityData;
import net.minecraft.entity.EntityType;
import net.minecraft.entity.SpawnReason;
import net.minecraft.entity.mob.DrownedEntity;
import net.minecraft.entity.mob.ZombieEntity;
import net.minecraft.entity.projectile.TridentEntity;
import net.minecraft.item.ItemStack;
import net.minecraft.item.Items;
import net.minecraft.nbt.NbtCompound;
import net.minecraft.nbt.NbtElement;
import net.minecraft.util.Identifier;
import net.minecraft.world.LocalDifficulty;
import net.minecraft.world.ServerWorldAccess;
import net.minecraft.world.World;
import org.jetbrains.annotations.Nullable;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.ModifyVariable;
import stone926.mods.more_enchantments.interfaces.IEntitySpawnedByFlower;

import java.util.Iterator;

@Mixin(DrownedEntity.class)
public abstract class DrownedEntityMixin extends ZombieEntity {

  @Shadow
  public abstract EntityData initialize(ServerWorldAccess world, LocalDifficulty difficulty, SpawnReason spawnReason, @Nullable EntityData entityData, @Nullable NbtCompound entityNbt);

  protected DrownedEntityMixin(EntityType<? extends ZombieEntity> entityType, World world) {
    super(entityType, world);
  }


  @ModifyVariable(method = "attack", at = @At(value = "STORE", ordinal = 0), ordinal = 0)
  private TridentEntity attack(TridentEntity trident) {
    ItemStack stack = getMainHandStack().copy();
    if (((IEntitySpawnedByFlower) this).isSpawnedByMandala() && getMainHandStack().isOf(Items.TRIDENT)) {
      Iterator<NbtElement> iterator = stack.getEnchantments().iterator();
      while (iterator.hasNext()) {
        NbtElement nbtElement = iterator.next();
        if (nbtElement instanceof NbtCompound compound) {
          Identifier id = EnchantmentHelper.getIdFromNbt(compound);
          if (EnchantmentHelper.getEnchantmentId(Enchantments.LOYALTY).equals(id)) {
            iterator.remove();
          }
        }
      }
      return new TridentEntity(world, this, stack);
    } else return trident;
  }

}
